__attribute__((global, cooperative)) void kernel_sdpa(
    float *__restrict dev_query, float *__restrict dev_key,
    float *__restrict dev_value, float *__restrict dev_out, const int bs,
    const int head_n, const int seq_len, const int head_size)
{
    
    
    // 获得线程数量
    int thread_num = GetThreadNum();
    // 获得当前线程ID
    int thread_idx = GetThreadIdx();

    printf("%d", head_n);
}

void GCU_SDPA(float *__restrict dev_query, float *__restrict dev_key,
              float *__restrict dev_value, float *__restrict dev_out,
              const int bs, const int head_n, const int seq_len,
              const int head_size)
{
    // 设置kernel参数
    static const size_t blocks = 1;
    static const size_t threads = 1;

    // 调用kernel
    kernel_sdpa<<<blocks, threads>>>(dev_query, dev_key, dev_value, dev_out, bs,
                                     head_n, seq_len, head_size);
}